import os
import random
import sys
from enum import Enum
import psutil
import shutil

CLUS_SIZE = 0x20
ROOT_START_IDX = 16  # 以M为单位
RESERVE_SIZE = 0x5000  # 保留扇区数 10M/512

class SPEED_FILL(Enum):
    FAST = 1 # 只将第4个簇填充为未使用
    MID = 2 # 只将前100M或小于100M的FAT填充为未使用
    SLOW = 3 # 将整个FAT填充为未使用


def fill_fat(fd, fat_size: int, tp: SPEED_FILL):
    
    data = int(0x0ffffff8).to_bytes(4, 'little')
    data = data + int(0xffffffff).to_bytes(4, 'little')
    data = data + int(0x0fffffff).to_bytes(4, 'little')
    if tp==SPEED_FILL.FAST:
        data = data+int(0).to_bytes(4,'little')
        os.write(fd, data)
        return
    data = data + (b'\x00' * 500)
    os.write(fd,data)
    fat_size -= 1
    i_count=0

    while True:
        if fat_size <= 0:
            break
        
        tmp_size = fat_size - 2048*10
        if (tmp_size < 0):
            tmp_size = fat_size
        else:
            tmp_size=2048*10
        data = b'\x00' * 512 * tmp_size
        os.write(fd, data)
        fat_size -= tmp_size
        i_count+=1
        # if (i_count%10==0):
        print('已初始化{}M'.format(i_count*10))

def fill_root_dir(fd):
    data = b'HWCH    ' # 文件名
    data = data + b'   ' # 扩展名
    data = data + b'\x08' #属性字节
    data = data + b'\x00' #系统保留
    data = data + b'\x00' # 创建时间的10毫秒
    data = data + b'\x00\x00' # 文件创建时间
    data = data + b'\x00\x00' #文件创建日期
    data = data + b'\x00\x00' # 文件最后访问时间
    data = data + b'\x00\x00' # 文件起始簇号高16位
    data = data + b'\x00\x00' # 文件最近修改时间
    data = data + b'\x00\x00' # 文件最近修改日期
    data = data + b'\x00\x00' # 文件起始簇号低16位
    data = data + b'\x00\x00\x00\x00' #文件长度
    os.write(fd, data)
    size = CLUS_SIZE * 512 - 32
    data = b'\x00' * size
    os.write(fd, data)

def get_mount_point(diskpath):
    for part in psutil.disk_partitions():
        if part.device.find(diskpath) != -1:
            return part.mountpoint

    return None

def write_image(disk: str, mbr: str, boot: str, kernel: str):
    

    mfd = os.open(mbr, os.O_RDONLY)
    bfd = os.open(boot, os.O_RDONLY)
    dfd = os.open(disk, os.O_RDWR)
    # ===============MBR===============
    pos=os.lseek(dfd,90,os.SEEK_CUR)
    print('开始写入MBR')
    data=os.read(mfd, 422)
    if len(data)<422:
        print('MBR镜像镜像不完整')
        exit(-1)
    else:
        os.write(dfd, data)
        print('写入MBR完成')

    # ===============Bootlaoder===============
    pos=os.lseek(dfd,16*512,os.SEEK_CUR)
    print('开始写入Bootloader')
    data=os.read(bfd, 512*1024)
    if len(data)<=0:
        print('Bootloader镜像镜像不完整')
        exit(-1)
    else:
        os.write(dfd, data)
        print('写入Bootloader完成')


    
    
    os.close(mfd)
    os.close(bfd)
    

    # ===============Kernel===============
    print('开始写入Kernel')
    d_point=get_mount_point(disk)
    if d_point is None:
        print('磁盘设备未挂在')
        exit(-1)
    else:
        shutil.copy2(kernel, d_point)
        print('写入Kernel完成')

    os.fsync(dfd)
    os.close(dfd)
    print('全部完成！！！')


if __name__ == '__main__':
    mbr=None
    boot=None
    disk=None
    kernel=None
    for arg in sys.argv[1:]:
        if arg.startswith('--mbr=') and len(arg)>6:
            mbr=arg[6:]
        elif arg.startswith('--boot=') and len(arg)>7:
            boot=arg[7:]
        elif arg.startswith('--disk=') and len(arg)>7:
            disk=arg[7:]
        elif arg.startswith('--kernel=') and len(arg)>9:
            kernel=arg[9:]
        elif arg=='--help':
            print('{} --mbr=<mbr image> --boot=<bootloader image> --disk=<disk device> --kernel=<kernel image>'.format(sys.argv[0]))
            exit(0)

    if mbr is None:
        print('MBR镜像未指定')
        exit(-1)

    if boot is None:
        print('Bootloader镜像未指定')
        exit(-1)

    if disk is None:
        print('磁盘设备未指定')
        exit(-1)

    if kernel is None:
        print('内核未指定')
        exit(-1)

    write_image(kernel=kernel, disk=disk, mbr=mbr, boot=boot)
    
    
    # write_dos(sys.argv[1], SPEED_FILL.MID)
